import torch
import torch.nn as nn
import torch.nn.functional as F

class DCTLoss(nn.Module):
    def __init__(self):
        super(DCTLoss, self).__init__()
        self.mse = nn.MSELoss()

    def forward(self, pred, target):
        # pred, target: (B, C, H, W)
        pred_dct = self.dct_2d(pred)
        target_dct = self.dct_2d(target)
        return self.mse(pred_dct, target_dct)

    def dct_2d(self, img):
        # Apply DCT along H and W
        return self.dct(self.dct(img, norm='ortho'), dim=-2, norm='ortho')

    def dct(self, x, dim=-1, norm=None):
        # 1D DCT along specified dim
        N = x.size(dim)
        v = torch.cat([x, x.flip([dim])], dim=dim)
        Vc = torch.fft.fft(v, dim=dim)
        k = torch.arange(N, device=x.device).view([1]*dim + [-1] + [1]*(x.dim()-dim-1))
        real = Vc.real.index_select(dim, torch.arange(N, device=x.device))  # keep only first N
        return real
